from flask import Flask, request, jsonify
from flask_cors import CORS
import requests
import json
import base64
import logging
import sys
import os
import time
from typing import List
from bs4 import BeautifulSoup
import urllib.parse
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from webdriver_manager.chrome import ChromeDriverManager
import atexit

# 配置日志
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('server.log')
    ]
)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)  # 使用最简单的 CORS 配置

API_KEY = "GtsPraOHhrwLeiby0ZaLLjLl"
SECRET_KEY = "1kgU9Pcz9UWs4Gjaw819dHs8PWj7hzD9"

# 修改 chrome_options 配置
chrome_options = Options()
chrome_options.add_argument('--headless')
chrome_options.add_argument('--no-sandbox')
chrome_options.add_argument('--disable-dev-shm-usage')
chrome_options.add_argument('--disable-gpu')
# 添加新的参数来处理 Ubuntu 服务器环境
chrome_options.add_argument('--disable-software-rasterizer')
chrome_options.add_argument('--disable-extensions')
chrome_options.add_argument('--remote-debugging-port=9222')

driver = None

def init_driver():
    global driver
    if driver is None:
        try:
            # 首先尝试安装必要的系统依赖
            os.system('apt-get update && apt-get install -y wget unzip chromium-browser')
            
            # 使用 ChromeDriverManager 安装驱动
            service = Service(ChromeDriverManager(chrome_type="chromium").install())
            driver = webdriver.Chrome(service=service, options=chrome_options)
            
            # 执行 JavaScript 来隐藏 WebDriver
            driver.execute_cdp_cmd('Page.addScriptToEvaluateOnNewDocument', {
                'source': '''
                    Object.defineProperty(navigator, 'webdriver', {
                        get: () => undefined
                    });
                    Object.defineProperty(navigator, 'plugins', {
                        get: () => [1, 2, 3, 4, 5]
                    });
                    window.chrome = {
                        runtime: {}
                    };
                '''
            })
            
        except Exception as e:
            logger.error(f"WebDriver 初始化失败: {str(e)}")
            raise
    return driver

def get_access_token():
    try:
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {
            "grant_type": "client_credentials",
            "client_id": API_KEY,
            "client_secret": SECRET_KEY
        }
        logger.info("正在获取访问令牌...")
        response = requests.post(url, params=params)
        response.raise_for_status()
        token = response.json().get("access_token")
        if not token:
            raise ValueError("未能获取到访问令牌")
        logger.info("成功获取访问令牌")
        return token
    except Exception as e:
        logger.error(f"获取访问令牌失败: {str(e)}")
        raise

@app.route('/api/image-understanding/submit', methods=['POST', 'OPTIONS'])
def submit_image_understanding():
    if request.method == 'OPTIONS':
        return '', 204
        
    try:
        logger.info("接收到图像理解请求")
        
        if not request.is_json:
            logger.error("请求不是JSON格式")
            return jsonify({"error": "请求必须是JSON格式"}), 400
            
        if not request.json:
            logger.error("未接收到JSON数据")
            return jsonify({"error": "未接收到JSON数据"}), 400
        
        image_data = request.json.get('image')
        if not image_data:
            logger.error("未接收到图像数据")
            return jsonify({"error": "未接收到图像数据"}), 400
        
        access_token = get_access_token()
        
        # 提交请求
        submit_url = f"https://aip.baidubce.com/rest/2.0/image-classify/v1/image-understanding/request?access_token={access_token}"
        
        # 准备form-data格式的数据
        submit_data = {
            "image": image_data,
            "question": "给定图像理解其所表达的含义和中心思想功能"
        }
        
        logger.info("正在提交图像理解请求...")
        try:
            submit_response = requests.post(
                submit_url,
                headers={'Content-Type': 'application/x-www-form-urlencoded'},
                data=submit_data
            )
            
            # 记录原始响应内容，用于调试
            logger.debug(f"原始响应内容: {submit_response.text}")
            
            submit_response.raise_for_status()
            submit_result = submit_response.json()
            logger.debug(f"提交响应: {submit_result}")
            
            if 'error_code' in submit_result:
                error_msg = submit_result.get('error_msg', '未知错误')
                logger.error(f"API返回错误: {error_msg}")
                return jsonify({
                    "status": "FAILED",
                    "result": f"API错误: {error_msg}"
                })
                
        except requests.exceptions.RequestException as e:
            logger.error(f"提交请求异常: {str(e)}")
            return jsonify({
                "status": "FAILED",
                "result": f"请求失败: {str(e)}"
            }), 500
        except json.JSONDecodeError as e:
            logger.error(f"解析响应JSON失败: {str(e)}, 原始响应: {submit_response.text}")
            return jsonify({
                "status": "FAILED",
                "result": "解析响应失败"
            }), 500
        
        # 获取task_id
        task_id = None
        if isinstance(submit_result.get('result'), dict):
            task_id = submit_result['result'].get('task_id')
        elif 'task_id' in submit_result:
            task_id = submit_result['task_id']
            
        if not task_id:
            logger.error("未获取到task_id")
            return jsonify({
                "status": "FAILED",
                "result": "未获取到任务ID"
            })
            
        logger.info(f"获取到task_id: {task_id}")
        
        # 轮询结果
        max_retries = 10
        retry_interval = 5
        
        for retry_count in range(max_retries):
            logger.info(f"第 {retry_count + 1} 次获取结果")
            time.sleep(retry_interval)
            
            result_url = f"https://aip.baidubce.com/rest/2.0/image-classify/v1/image-understanding/get-result?access_token={access_token}"
            result_data = json.dumps({
                "task_id": str(task_id)
            }, ensure_ascii=False)
            
            result_response = requests.post(
                result_url,
                headers={'Content-Type': 'application/json'},
                data=result_data.encode('utf-8')
            )
            result_response.raise_for_status()
            result = result_response.json()
            logger.debug(f"获取结果响应: {result}")
            
            # 获取嵌套的结果数据
            result_data = result.get('result', {})
            ret_code = result_data.get('ret_code')
            
            if ret_code == 0:
                # 成功获取结果
                description = result_data.get('description')
                answer = result_data.get('answer')
                
                if description:
                    logger.info(f"获取到结果(description): {description}")
                    return jsonify({
                        "status": "FINISHED",
                        "result": description
                    })
                elif answer:
                    logger.info(f"获取到结果(answer): {answer}")
                    return jsonify({
                        "status": "FINISHED",
                        "result": answer
                    })
                else:
                    logger.warning("结果中既没有description也没有answer字段")
                    if retry_count == max_retries - 1:
                        return jsonify({
                            "status": "FAILED",
                            "result": "无法获取有效的分析结果"
                        })
                    continue
            elif ret_code == 1:
                # 任务仍在处理中
                logger.info("任务仍在处理中，继续等待...")
                if retry_count == max_retries - 1:
                    logger.warning("达到最大重试次数，返回超时")
                    return jsonify({
                        "status": "TIMEOUT",
                        "result": "获取结果超时"
                    })
                continue
            else:
                # 其他错误情况
                error_msg = result_data.get('ret_msg', '获取结果失败')
                logger.error(f"获取结果失败: {error_msg}")
                return jsonify({
                    "status": "FAILED",
                    "result": error_msg
                })
        
    except Exception as e:
        logger.error(f"处理请求失败: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/video-search/submit', methods=['POST', 'OPTIONS'])
def search_video():
    """
    使用 Selenium 搜索B站视频
    """
    if request.method == 'OPTIONS':
        return '', 204
    
    try:
        if not request.is_json:
            return jsonify({"error": "请求必须是JSON格式"}), 400

        search_text = request.json.get('query')
        if not search_text:
            return jsonify({"error": "搜索文本不能为空"}), 400
        
        logger.info(f"开始搜索视频，关键词: {search_text}")
        
        try:
            # 初始化 driver
            driver = init_driver()
            logger.info("WebDriver 初始化成功")
            
            # 访问B站搜索页面
            search_url = f"https://search.bilibili.com/all?keyword={urllib.parse.quote(search_text)}"
            logger.info(f"访问搜索URL: {search_url}")
            driver.get(search_url)
            
            # 等待页面加载
            time.sleep(3)
            
            # 等待视频元素加载
            wait = WebDriverWait(driver, 15)
            logger.info("等待视频元素加载...")
            
            try:
                video_items = wait.until(
                    EC.presence_of_all_elements_located(
                        (By.CSS_SELECTOR, ".video-item")
                    )
                )
                logger.info(f"找到 {len(video_items)} 个视频")
            except Exception as e:
                logger.error(f"等待视频元素超时: {str(e)}")
                video_items = driver.find_elements(By.CSS_SELECTOR, ".bili-video-card,.video-item,.search-video-card")
                logger.info(f"使用备选选择器找到 {len(video_items)} 个视频")
            
            video_results = []
            
            # 获取视频结果
            for item in video_items[:10]:
                try:
                    # 只获取标题、链接和时长
                    title = ""
                    try:
                        title = item.find_element(By.CSS_SELECTOR, ".bili-video-card__info--tit,.title,.video-title").text
                    except:
                        title = item.find_element(By.CSS_SELECTOR, "a[title]").get_attribute("title")
                    
                    url = item.find_element(By.CSS_SELECTOR, "a[href*='video']").get_attribute("href")
                    
                    # 获取时长
                    duration = ""
                    try:
                        duration = item.find_element(By.CSS_SELECTOR, ".bili-video-card__stats__duration,.duration,.video-duration").text
                    except:
                        duration = "未知时长"
                    
                    video_results.append({
                        'title': title,
                        'url': url,
                        'duration': duration
                    })
                    
                    logger.debug(f"成功解析视频: {title}")
                    
                except Exception as e:
                    logger.error(f"解析单个视频项时出错: {str(e)}")
                    continue
            
            if not video_results:
                logger.warning("未找到任何视频结果")
                return jsonify({
                    "status": "success",
                    "results": []
                })
            
            logger.info(f"成功获取 {len(video_results)} 个视频结果")
            return jsonify({
                "status": "success",
                "results": video_results
            })
            
        except Exception as e:
            logger.error(f"Selenium 操作失败: {str(e)}")
            raise

    except Exception as e:
        error_msg = f"视频搜索失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

@app.route('/api/image-retrieval/submit', methods=['POST', 'OPTIONS'])
def search_similar_images():
    """
    处理图片上传并返回
    """
    if request.method == 'OPTIONS':
        return '', 204
    
    try:
        if not request.is_json:
            return jsonify({"error": "请求必须是JSON格式"}), 400
        
        image_data = request.json.get('image')
        if not image_data:
            return jsonify({"error": "图片数据不能为空"}), 400
        
        logger.info("接收到图片数据")
        
        # 直接返回上传的图片数据
        image_results = [{
            'url': f"data:image/jpeg;base64,{image_data}"
        }]
        
        return jsonify({
            "status": "success",
            "results": image_results
        })
            
    except Exception as e:
        error_msg = f"图片处理失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

# 在应用退出时关闭driver
@atexit.register
def cleanup():
    global driver
    if driver:
        driver.quit()

# 确保在应用启动时初始化 driver
try:
    init_driver()
    logger.info("WebDriver 初始化成功")
except Exception as e:
    logger.error(f"WebDriver 初始化失败: {str(e)}")

if __name__ == '__main__':
    try:
        logger.info("正在启动服务器...")
        app.run(host='0.0.0.0', port=5003, debug=True)
    except Exception as e:
        logger.error(f"服务器启动失败: {str(e)}")
        sys.exit(1) 